{-# LANGUAGE PackageImports #-}

module Env where

import Data.Map (Map)
import qualified Data.Map as Map
import "mtl" Control.Monad.State.Strict
import Types

type NameGenerator = Integer
    
newName :: NameGenerator -> String -> (NameGenerator, String)
newName g x = (g + 1, x ++ "_" ++ show g)


-- LEnv is a mapping from local variables to types
type LEnv = Map String Type

-- FEnv is a mapping from function names to function types
type FEnv = Map String FunType

-- UEnv is a mapping from function names to user-defined functions
type UEnv = Map String Function

-- Env is a record storing all environments
data Env = Env {
    lenv :: LEnv,
    fenv :: FEnv,
    uenv :: UEnv,           -- untranslated user-defined functions
    tenv :: UEnv,           -- translated user-defined functions
    gen  :: NameGenerator,
    desiredRetType :: Type  -- enforce return type of function being translated
}

-- the translator monad is a state monad holding an Env
type Trans a = State Env a

-- initial environment
initEnv :: Env
initEnv = Env Map.empty builtinFunc Map.empty Map.empty 1 NoType

-- evaluation function for Trans
evalTrans :: Trans a -> a
evalTrans t = evalState t initEnv


-- the types of builtin functions
builtinFunc ::  FEnv
builtinFunc =
    Map.fromList [
        ("lt",          FunType [Float, Float] Boolean),
        ("gt",          FunType [Float, Float] Boolean),
        ("eq",          FunType [Float, Float] Boolean),
        ("and",         FunType [Boolean, Boolean] Boolean),
        ("or",          FunType [Boolean, Boolean] Boolean),
        ("not",         FunType [Boolean] Boolean),

        ("data1d_rgba", FunType [String, String] Dataset),
        ("data1d",      FunType [String, String] Dataset),
        ("data2d_rgba", FunType [String, String] Dataset),
        ("data2d",      FunType [String, String] Dataset),
        ("data3d",      FunType [String, String] Dataset),
        ("data4d",      FunType [String, String] Dataset),

        ("float_param",  FunType [Float, Float, Float, String] Float),
        ("float3_param", FunType [Float, Float, Float] Float3),

        ("make_float2", FunType [Float, Float] Float2),
        ("make_float3", FunType [Float, Float, Float] Float3),
        ("make_float4", FunType [Float, Float, Float, Float] Float4),

        ("make_int2", FunType [Int, Int] Int2),
        ("make_int3", FunType [Int, Int, Int] Int3),
        ("make_int4", FunType [Int, Int, Int, Int] Int4),

        ("cubic_query_3d",     FunType [Dataset, Float3] Float),
        ("cubic_gradient_3d",  FunType [Dataset, Float3] Float3),
        ("linear_query_3d",    FunType [Dataset, Float3] Float),
        ("linear_gradient_3d", FunType [Dataset, Float3] Float3),

        ("cubic_query_3d_cut",     FunType [Dataset, Float3, Float3, Float] Float),
        ("cubic_gradient_3d_cut",  FunType [Dataset, Float3, Float3, Float] Float3),
        ("linear_query_3d_cut",    FunType [Dataset, Float3, Float3, Float] Float),
        ("linear_gradient_3d_cut", FunType [Dataset, Float3, Float3, Float] Float3),

        ("cubic_query_4d",     FunType [Dataset, Float3, Float] Float),
        ("cubic_gradient_4d",  FunType [Dataset, Float3, Float] Float3),
        ("linear_query_4d",    FunType [Dataset, Float3, Float] Float),
        ("linear_gradient_4d", FunType [Dataset, Float3, Float] Float3),

        ("cubic_query_4d_cut",     FunType [Dataset, Float3, Float3, Float, Float] Float),
        ("cubic_gradient_4d_cut",  FunType [Dataset, Float3, Float3, Float, Float] Float3),
        ("linear_query_4d_cut",    FunType [Dataset, Float3, Float3, Float, Float] Float),
        ("linear_gradient_4d_cut", FunType [Dataset, Float3, Float3, Float, Float] Float3),

        ("linear_query_2d_rgba", FunType [Dataset, Float2] Float4),
        ("linear_query_2d",      FunType [Dataset, Float2] Float),
        ("linear_query_1d_rgba", FunType [Dataset, Float] Float4),
        ("linear_query_1d",      FunType [Dataset, Float] Float),

        ("dot",         FunType [Float3, Float3] Float),
        ("normalize",   FunType [Float3] Float3),
        ("length",      FunType [Float3] Float),
        ("lengthSqr",   FunType [Float3] Float),
        ("cross",       FunType [Float3, Float3] Float3),
        
        ("phong",       FunType [Float3, Float3, Float3, Float3, Float3, Float, Float] Float3),

        ("sqrt",        FunType [Float] Float),
        ("pow",         FunType [Float, Float] Float),
		("exp",         FunType [Float] Float),
        ("abs",         FunType [Float] Float),
        
        ("sin",         FunType [Float] Float),
        ("cos",         FunType [Float] Float),
        ("asin",        FunType [Float] Float),
        ("acos",        FunType [Float] Float),
        ("atan",        FunType [Float] Float),
        ("atan2",       FunType [Float, Float] Float),
        
        ("clamp",       FunType [Float, Float, Float] Float),
        ("clamp3",      FunType [Float3, Float3, Float3] Float3),
        
        ("step",        FunType [Float, Float] Float),
        ("smoothstep",  FunType [Float, Float, Float] Float),
        
		("square",      FunType [Float] Float),
        ("floor",       FunType [Float] Float),
        ("ceil",        FunType [Float] Float),
        ("frac",        FunType [Float] Float),
        ("rand",        FunType [] Float),
        ("fmod",        FunType [Float, Float] Float)
    ]

